
#include "RainbowTableGenerator.h"
#include "ChainWalkContext.h"
#include <iostream>
#include <sstream>
#include <time.h>
#include <conio.h>

CRainbowTableGenerator::CRainbowTableGenerator()
{
	m_nCurrentCalculatedChains = 0;
}

CRainbowTableGenerator::~CRainbowTableGenerator(void)
{
}

int CRainbowTableGenerator::CalculateTable(std::string sFilename, int nRainbowChainCount, std::string sHashRoutineName, std::string sCharsetName, int nPlainLenMin, int nPlainLenMax, int nRainbowTableIndex, int nRainbowChainLen, uint64 nChainStart, std::string sSalt)
{
	// Setup CChainWalkContext
	//std::cout << "ChainGenerator started." << std::endl;

	if (!CChainWalkContext::SetHashRoutine(sHashRoutineName))
	{
		std::cerr << "hash routine " << sHashRoutineName << " not supported" << std::endl;
		return 1;
	}
	//std::cout << "Hash routine validated" << std::endl;

	if (!CChainWalkContext::SetPlainCharset(sCharsetName, nPlainLenMin, nPlainLenMax))
	{	
		std::cerr << "charset " << sCharsetName << " not supported" << std::endl;
		return 2;
	}
	//std::cout << "Plain charset validated" << std::endl;

	if (!CChainWalkContext::SetRainbowTableIndex(nRainbowTableIndex))
	{
		std::cerr << "invalid rainbow table index " << nRainbowTableIndex << std::endl;
		return 3;
	}
	//std::cout << "Rainbowtable index validated" << std::endl;

	if(sHashRoutineName == "mscache")// || sHashRoutineName == "lmchall" || sHashRoutineName == "halflmchall")
	{
		// Convert username to unicode
		const char *szSalt = sSalt.c_str();
		int salt_length = strlen(szSalt);
		unsigned char cur_salt[256];
		for (int i=0; i<salt_length; i++)
		{
			cur_salt[i*2] = szSalt[i];
			cur_salt[i*2+1] = 0x00;
		}
		CChainWalkContext::SetSalt(cur_salt, salt_length*2);
	}
	else if(sHashRoutineName == "halflmchall")
	{ // The salt is hardcoded into the hash routine
	//	CChainWalkContext::SetSalt((unsigned char*)&salt, 8);
	}
	else if(sHashRoutineName == "oracle")
	{
		CChainWalkContext::SetSalt((unsigned char *)sSalt.c_str(), sSalt.length());
	}
	//std::cout << "Opening chain file" << std::endl;

	
	// Open file
	fclose(fopen(sFilename.c_str(), "a"));
	FILE* file = fopen(sFilename.c_str(), "r+b");
	
	if (file == NULL)
	{
		std::cerr << "failed to create " << sFilename << std::endl;
		return 4;
	}
	
	
	// Check existing chains
	unsigned int nDataLen = (unsigned int)GetFileLen(file);
//	unsigned int nFileLen;
	
	nDataLen = nDataLen / 16 * 16;
	if ((int)nDataLen == nRainbowChainCount * 16)
	{		
		std::cerr << "precomputation of this rainbow table already finished" << std::endl;
		fclose(file);
		return 0;
	}
	nChainStart += (nDataLen / 16);
	fseek(file, nDataLen, SEEK_SET);
	size_t nReturn;
	CChainWalkContext cwc;
	uint64 nIndex[2];
	time_t tStart = time(NULL);
//	std::cout << "Starting to generate chains" << std::endl;

	for(m_nCurrentCalculatedChains = nDataLen / 16; m_nCurrentCalculatedChains < nRainbowChainCount; m_nCurrentCalculatedChains++)
	{		
		//cwc.GenerateRandomIndex();
		cwc.SetIndex(nChainStart++); // use a given index now!
		nIndex[0] = cwc.GetIndex();
		for (int nPos = 0; nPos < nRainbowChainLen - 1; nPos++)
		{
		//	std::cout << "IndexToPlain()" << std::endl;
			cwc.IndexToPlain();
		//	std::cout << "PlainToHash()" << std::endl;
			cwc.PlainToHash();
		//	std::cout << "HashToIndex()" << std::endl;
			cwc.HashToIndex(nPos);
		}
		//std::cout << "GetIndex()" << std::endl;

		nIndex[1] = cwc.GetIndex();
		// Write chain to disk
		if ((nReturn = fwrite(nIndex, 1, 16, file)) != 16)
		{
			std::cerr << "disk write fail" << std::endl;
			fclose(file);
			return 9;
		}
		fflush(file);
	}
	//std::cout << "Generation completed" << std::endl;

	fclose(file);
	return 0;
}

